import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(1, 50), nn.ReLU(),
            nn.Linear(50, 100), nn.ReLU(),
            nn.Linear(100, 50), nn.ReLU(),
            nn.Linear(50, 1)
        )
    def forward(self, x):
        return self.net(x)

a = torch.linspace(-7, 7, 400)
b = torch.sin(a)

A = a.unsqueeze(1)
B = b.unsqueeze(1)

data = list(zip(A, B))

net = MyNet()

batchsize = 20
dataloader = DataLoader(data, batch_size=batchsize, shuffle=True)
# 损失函数
loss_fun = nn.MSELoss()
# 优化器
optimer = torch.optim.SGD(net.parameters(), lr = 0.01)

epochnum = 600

ls = []
for epoch in range(epochnum):
    train_loss = 0.0
    for x, y in dataloader:
        optimer.zero_grad()
        train_x = net.forward(x)
        loss = loss_fun(train_x, y)
        loss.backward()
        optimer.step()
        train_loss += loss.item()
    ls.append(train_loss/batchsize)
    if (epoch+1)%10 == 0:
        print(f'epoch:{epoch+1}, loss:{train_loss/batchsize:.7f}')

plt.plot(a.numpy(), net(A).detach().numpy())
plt.plot(a.numpy(), b.numpy(), c='r')
plt.show()
plt.plot(range(epochnum), ls)
plt.show()
